/*
 * Copyright 2006-2007 S W McIlroy
 * Licensed  under the  Apache License,  Version 2.0  (the "License");
 * you may not use  this file  except in  compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed  under the  License is distributed on an "AS IS" BASIS,
 * WITHOUT  WARRANTIES OR CONDITIONS  OF ANY KIND, either  express  or
 * implied.
 *
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.mockfu;

import static org.mockfu.Exceptions.*;
import static org.objectweb.asm.Opcodes.*;

import java.lang.reflect.Constructor;
import java.lang.reflect.Field;
import java.lang.reflect.Method;
import java.lang.reflect.Modifier;
import java.lang.reflect.ParameterizedType;
import java.util.ArrayList;
import java.util.List;
import java.util.Set;
import java.util.TreeSet;

import org.objectweb.asm.ClassAdapter;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.FieldVisitor;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.Type;

/**
 * A <tt>UnitBuilder</tt> creates an object isolated from its dependencies through a set of
 * replacement operations applied to the bytecode of a class. The builder works by creating a
 * subclass of the unit which is identical to its superclass except for the bytecode modifications
 * made by each replacement operation.
 *
 * @param <U> the type of the unit class
 */
public class UnitBuilder<U>
{
    private static final String FINAL_METHOD_SUFFIX = "$Unit$internal";

    private Class<U> klas;
    private String klasRef;
    private String subKlasName;
    private String subKlasRef;

    private final List<Replacement<U>> replacements = new ArrayList<Replacement<U>>();
    private final Set<String> privateInnerClasses = new TreeSet<String>();

    /**
     * Constructs a builder for generating a unit of the specified type.
     *
     * @param type the class of the unit
     * @return the unit builder
     */
    public static <U> UnitBuilder<U> of(Class<U> type)
    {
        return new UnitBuilder<U>(type);
    }

    /**
     * Constructs a builder for generating a unit of the specified type.
     *
     * @param type the class of the unit
     */
    public UnitBuilder(Class<U> type)
    {
        setClass(type);
    }

    /**
     * Allows subclasses with generic type parameters to create a control object for a mock with
     * that type. For example, the following declaration creates a control for a mock implementing
     * the <tt>List</tt> interface without causing a raw type compiler warning:
     *
     * <pre>
     * MockControl&lt;List&lt;String&gt;&gt; list = new MockControl&lt;List&lt;String&gt;&gt;() {};
     * </pre>
     */
    @SuppressWarnings("unchecked")
    protected UnitBuilder()
    {
        java.lang.reflect.Type type = this.getClass().getGenericSuperclass();
        if (type instanceof Class)
            throw new RuntimeException("Missing type parameter");
        type = ((ParameterizedType)type).getActualTypeArguments()[0];
        if (type instanceof ParameterizedType)
            setClass((Class<U>)((ParameterizedType)type).getRawType());
        else
            setClass((Class<U>)type);
    }

    private void setClass(Class<U> klas)
    {
        if (!Modifier.isPublic(klas.getModifiers()) && !Modifier.isProtected(klas.getModifiers()))
            throw new IllegalArgumentException("Unit class must be either public or protected");
        if (Modifier.isFinal(klas.getModifiers()))
            throw new IllegalArgumentException("Unit class cannot be final");

        this.klas = klas;
        this.klasRef = klas.getName().replace(".", "/");
        this.subKlasName = klas.getName() + "$$Unit$$";
        this.subKlasRef = subKlasName.replace(".", "/");
    }

    public <C> MethodReplacement<U, C> replaceConstructor(Class<C> ctorClass, Object... args)
    {
        MethodReplacement<U, C> rep = new MethodReplacement<U, C>(klas, ctorClass, args);
        replacements.add(rep);
        return rep;
    }

    public MethodReplacement<U, Object> replaceInvocation(ConstrainedMethod method)
    {
        MethodReplacement<U, Object> rep = new MethodReplacement<U, Object>(klas, method);
        replacements.add(rep);
        return rep;
    }

    String getClassRef()
    {
        return klasRef;
    }

    String getSubclassName()
    {
        return subKlasName;
    }

    String getSubclassRef()
    {
        return subKlasRef;
    }

    @SuppressWarnings("unchecked")
    public U createUnit(Class<?>[] parameterTypes, Object... args)
    {
        try
        {
            ClassReader classReader = new ClassReader(klas.getName());
            final ClassWriter classWriter = new ClassWriter(ClassWriter.COMPUTE_MAXS);
            ClassVisitor visitor = classWriter;
//            visitor = new CheckClassAdapter(visitor);
//            visitor = new TraceClassVisitor(visitor, new PrintWriter(System.out));
            visitor = new UnitWriter(visitor);
            for (Replacement<U> replacement : replacements)
                visitor = replacement.chain(visitor);

            classReader.accept(visitor, ClassReader.SKIP_FRAMES);

            ClassLoader loader = new ClassLoader() {
                @Override
                protected Class<?> findClass(String name) throws ClassNotFoundException
                {
                    byte[] bytes = classWriter.toByteArray();
                    return super.defineClass(name, bytes, 0, bytes.length);
                }
            };

            Class<U> c = (Class<U>)loader.loadClass(getSubclassName());
            if (parameterTypes == null)
                parameterTypes = new Class<?>[0];

            Constructor<U> ctor = c.getDeclaredConstructor(parameterTypes);
            ctor.setAccessible(true);
            return ctor.newInstance(args);
        }
        catch (Exception ex)
        {
            throw undeclared(ex);
        }
    }

    public U createUnit(Object... args)
    {
        if (args == null)
            return createUnit(new Class<?>[0], new Object[0]);

        Class<?>[] parameterTypes = new Class[args.length];
        for (int i = 0; i < args.length; i++)
        {
            Object arg = args[i];
            parameterTypes[i] = arg == null ? Object.class : arg.getClass();
        }
        return createUnit(parameterTypes, args);
    }

    class UnitWriter extends ClassAdapter
    {

        public UnitWriter(ClassVisitor cw)
        {
            super(cw);
        }

        @Override
        public void visit(int version, int access, String name, String signature, String supername, String[] interfaces)
        {
            super.visit(version, access, getSubclassRef(), signature, getClassRef(), interfaces);
        }

        @Override
        public FieldVisitor visitField(int access, String name, String desc, String signature, Object value)
        {
            try
            {
                klas.getDeclaredField(name).setAccessible(true);
            }
            catch (Exception ex)
            {
                throw undeclared(ex);
            }
            return null;
        }

        @Override
        public void visitInnerClass(String name, String outerName, String innerName, int access)
        {
            if (!Modifier.isProtected(access) && !Modifier.isPublic(access))
            {
                privateInnerClasses.add(name);
            }

            return;
        }

        @Override
        public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions)
        {
            if (name.equals("<clinit>"))
            {
                /* skip static initializer since we're using static members of super class */
                return null;
            }

            if (name.equals("<init>"))
            {
                /* make constructor public so that we can invoke it from UnitBuilder */
                access = (access & -8) + 1;

                MethodVisitor mv = super.visitMethod(access, name, desc, signature, exceptions);
                /* just invoke super class constructor */
                Type[] params = Type.getArgumentTypes(desc);
                mv.visitCode();
                mv.visitVarInsn(Opcodes.ALOAD, 0);

                for (int i = 1 ; i <= params.length ; i++)
                {
                    Type param = params[i - 1];
                    if (param == Type.INT_TYPE || param == Type.SHORT_TYPE || param == Type.BYTE_TYPE
                            || param == Type.BOOLEAN_TYPE || param == Type.CHAR_TYPE)
                        mv.visitVarInsn(ILOAD, i);
                    else if (param == Type.LONG_TYPE)
                        mv.visitVarInsn(LLOAD, i);
                    else if (param == Type.FLOAT_TYPE)
                        mv.visitVarInsn(FLOAD, i);
                    else if (param == Type.DOUBLE_TYPE)
                        mv.visitVarInsn(DLOAD, i);
                    else
                        mv.visitVarInsn(ALOAD, i);
                }

                mv.visitMethodInsn(INVOKESPECIAL, getClassRef(), name, desc);
                mv.visitInsn(RETURN);
                mv.visitMaxs(0, 0);
                return null;
            }

            if (Modifier.isFinal(access))
            {
                /* append suffix to method name so it doesn't conflict with super class */
                name += FINAL_METHOD_SUFFIX;
            }

            MethodVisitor mv = super.visitMethod(access, name, desc, signature, exceptions);
            return new PrivateAccessVisitor(mv, access, name, desc);
        }

    }

    /**
     * Rewrites field accessor instructions for private variables to use reflection, since direct
     * access results an IllegalAccessError.
     */
    private class PrivateAccessVisitor extends NormalMethodGenerator
    {
        public PrivateAccessVisitor(MethodVisitor delegate, int access, String name, String desc)
        {
            super(delegate, access, name ,desc);
        }


        @Override
        public void checkCast(Type type)
        {
            if (type.getSort() == Type.ARRAY)
            {
                super.checkCast(type);
                return;
            }

            if (!privateInnerClasses.contains(type.getInternalName()))
            {
                super.checkCast(type);
                return;
            }

            try
            {
                /* cast to the first castable class in the type hierarchy */
                Class<?> klas = Class.forName(type.getClassName());

                for (Class<?> superClass = klas.getSuperclass(); superClass != Object.class; superClass = superClass.getSuperclass())
                {
                    Type superType = Type.getType(superClass);
                    if (!privateInnerClasses.contains(superType.getInternalName()))
                    {
                        super.checkCast(superType);
                        break;
                    }
                }
            }
            catch (ClassNotFoundException ex)
            {
                throw undeclared(ex);
            }
        }


        @Override
        public void visitTypeInsn(int opcode, String desc)
        {
            if (opcode == NEW)
            {
                if (!privateInnerClasses.isEmpty())
                {
                    /* use ConstructorReplacementVisitor on constructor calls, because we don't know
                     * until <init> is called whether any arguments are for private inner classes */
                    mv = new ConstructorReplacementVisitor(mv);
                }
                if (desc.equals(getClassRef()))
                {
                    /* create new instances of self using subclass */
                    desc = getSubclassRef();
                }
            }
            if (privateInnerClasses.contains(desc))
            {
                if (opcode == INSTANCEOF)
                {
                    /* rewrite instanceof instructions to use Class.isInstance() */
                    push(Type.getObjectType(desc).getClassName());
                    invokeStatic("java/lang/Class", "forName", "(Ljava/lang/String;)Ljava/lang/Class;");
                    swap();
                    invokeVirtual("java/lang/Class", "isInstance", "(Ljava/lang/Object;)Z");
                    return;
                }
                else if (opcode == CHECKCAST)
                {
                    /* rewrite cast instructions to use Class.checkcast() */
                    return;
                }
            }

            mv.visitTypeInsn(opcode, desc);
        }

        @Override
        public void visitLdcInsn(Object cst)
        {
            if (cst instanceof Type)
            {
                Type type = (Type)cst;
                if (privateInnerClasses.contains(type.getInternalName()))
                {
                    /* rewrite class literals to use Class.forName() */
                    push(type.getClassName());
                    invokeStatic("java/lang/Class", "forName", "(Ljava/lang/String;)Ljava/lang/Class;");
                    return;
                }
            }
            mv.visitLdcInsn(cst);
        }



        @Override
        public void visitMethodInsn(int opcode, String owner, String name, String desc)
        {
            try
            {

                if ((opcode == INVOKESPECIAL || opcode == INVOKESTATIC) && owner.equals(getClassRef()))
                {
                    /* call subclass method with same signature */
                    if (!name.equals("<init>"))
                    {
                        Method method = getMethod(owner, name, desc);
                        /* calls to final methods need to be munged */
                        if (Modifier.isFinal(method.getModifiers()))
                            name += FINAL_METHOD_SUFFIX;
                    }
                    owner = getSubclassRef();
                }

                boolean requiresReflection = requiresReflection(owner, name, desc);

                if (requiresReflection)
                {
                    if (name.equals("<init>"))
                    {
                        /* replace call to constructor of private inner class with Constructor.newInstance() */
                        reflectConstructor(owner, name, desc);
                    }
                    else
                    {
                        /* rewrite method calls on private inner classes to use reflection */
                        reflectMethod(owner, name, desc);
                    }

                    return;
                }
                else if (name.equals("<init>") && mv instanceof ConstructorReplacementVisitor)
                {
                    /* we set up a recorder but didn't require reflection after all: play back instructions */
                    ConstructorReplacementVisitor recorder = (ConstructorReplacementVisitor)mv;
                    mv = recorder.playback();
                }
            }
            catch (Exception ex)
            {
                throw undeclared(ex);
            }

            mv.visitMethodInsn(opcode, owner, name, desc);
        }


        private void reflectMethod(String owner, String name, String desc)
        {
            Method method = getMethod(owner, name, desc);
            Class<?> owningClass = method.getDeclaringClass();
            Type owningType = Type.getType(owningClass);


            int arraySlot = storeArray(Type.getArgumentTypes(desc));

            /* push class and method name */
            push(owningType.getClassName());
            invokeStatic("java/lang/Class", "forName", "(Ljava/lang/String;)Ljava/lang/Class;");
            push(name);

            /* create a class array with the method parameter types */
            Class<?>[] parameterTypes = method.getParameterTypes();
            pushClassArray(parameterTypes);

            invokeVirtual("java/lang/Class", "getDeclaredMethod",
                          "(Ljava/lang/String;[Ljava/lang/Class;)Ljava/lang/reflect/Method;");

            /* make method accessible so it can be invoked */
            dup();
            push(true);
            invokeVirtual("java/lang/reflect/Method", "setAccessible", "(Z)V");

            /* invoke method */
            if (Modifier.isStatic(method.getModifiers()))
                pushNull();
            else
                swap();
            loadLocal(arraySlot);
            invokeVirtual("java/lang/reflect/Method", "invoke",
                          "(Ljava/lang/Object;[Ljava/lang/Object;)Ljava/lang/Object;");
            /* cast or unbox return value */
            unbox(Type.getReturnType(method));
        }


        private void reflectConstructor(String owner, String name, String desc)
        {
            if (!(mv instanceof ConstructorReplacementVisitor))
                throw new IllegalStateException("call to <init> without new");

            /* stop recording and playback arguments */
            ConstructorReplacementVisitor ctorReplacer = (ConstructorReplacementVisitor)mv;
            mv = ctorReplacer.playbackArguments();

            /* create an array with constructor arguments from top of stack */
            int arraySlot = storeArray(Type.getArgumentTypes(desc));

            /* get constructor object */
            Type owningType = Type.getObjectType(owner);
            push(owningType.getClassName());
            invokeStatic("java/lang/Class", "forName", "(Ljava/lang/String;)Ljava/lang/Class;");
            pushClassArray(getClasses(Type.getArgumentTypes(desc)));
            invokeVirtual("java/lang/Class", "getDeclaredConstructor",
                          "([Ljava/lang/Class;)Ljava/lang/reflect/Constructor;");

            /* make constructor accessible so it can be invoked */
            dup();
            push(true);
            invokeVirtual("java/lang/reflect/Constructor", "setAccessible", "(Z)V");

            /* invoke constructor */
            loadLocal(arraySlot);
            invokeVirtual("java/lang/reflect/Constructor", "newInstance",
                          "([Ljava/lang/Object;)Ljava/lang/Object;");
            checkCast(owningType);

            ctorReplacer.adjustCopies();
        }


        /**
         * If owner or any argument is a private inner class, method needs to be rewritten to use
         * reflection.
         */
        private boolean requiresReflection(String owner, String name, String desc)
        {
            if (privateInnerClasses.contains(owner))
                return true;

            Class<?>[] params = getClasses(Type.getArgumentTypes(desc));
            for (Class<?> param : params)
            {
                Type parameterType = Type.getType(param);
                if (parameterType.getSort() == Type.ARRAY)
                    parameterType = parameterType.getElementType();

                if (parameterType.getSort() == Type.OBJECT
                        && privateInnerClasses.contains(parameterType.getInternalName()))
                    return true;
            }

            return false;
        }


        /**
         * Creates a literal class array on the top of the static (using reflection on private inner
         * class types).
         */
        @Override
        public void pushClassArray(Class<?>[] classes)
        {
            push(classes.length);
            newArray(Type.getType(Class.class));

            for (int i = 0 ; i < classes.length ; i++)
            {
                dup();
                push(i);

                Type parameterType = Type.getType(classes[i]);
                if (parameterType.getSort() == Type.ARRAY)
                    parameterType = parameterType.getElementType();

                if (parameterType.getSort() == Type.OBJECT
                        && privateInnerClasses.contains(parameterType.getInternalName()))
                {
                    push(classes[i].getName());
                    invokeStatic("java/lang/Class", "forName", "(Ljava/lang/String;)Ljava/lang/Class;");
                }
                else
                {
                    pushBoxType(Type.getType(classes[i]));
                }
                arrayStore(Type.getType(Class.class));
            }
        }

        @Override
        public void visitFieldInsn(int opcode, String owner, String name, String desc)
        {
            if (owner.equals(getClassRef()) || privateInnerClasses.contains(owner))
            {
                /* rewrite private accessor instructions to use reflection */
                try
                {
                    Field field = getField(owner, name, desc);
                    Class<?> owningClass = field.getDeclaringClass();
                    Type owningType = Type.getType(owningClass);


                    if ((!Modifier.isPublic(field.getModifiers()) && !Modifier.isProtected(field.getModifiers()))
                            || privateInnerClasses.contains(owner))
                    {
                        boolean isCategory2 = field.getType() == long.class || field.getType() == double.class;

                        if (opcode == PUTFIELD && isCategory2)
                        {
                            /* reverse stack operands so they're easier to reorder later */
                            dup2X1();
                            pop2();
                        }

                        /* get field object */
                        if (privateInnerClasses.contains(owner))
                        {
                            /* can't load class explicitly, use Class.forName() */
                            push(owningType.getClassName());
                            invokeStatic("java/lang/Class", "forName", "(Ljava/lang/String;)Ljava/lang/Class;");
                        }
                        else
                            push(owningType);
                        push(name);
                        invokeVirtual("java/lang/Class", "getDeclaredField",
                                      "(Ljava/lang/String;)Ljava/lang/reflect/Field;");
                        /* make field accessible thru reflection */
                        dup();
                        push(true);
                        invokeVirtual("java/lang/reflect/Field", "setAccessible", "(Z)V");

                        /*
                         * rearrange the stack so that the field accessor has its arguments in the
                         * right order
                         */
                        String methodPrefix = "get";
                        if (opcode == GETFIELD)
                        {
                            /* swap field object and owning object */
                            swap();
                        }
                        else if (opcode == PUTFIELD)
                        {
                            methodPrefix = "set";
                            /* put field object before owning object and value to set */
                            if (isCategory2)
                            {
                                swap();
                                dup2X2();
                                pop();
                                pop();
                            }
                            else
                            {
                                dupX2();
                                pop();
                            }
                        }
                        else if (opcode == GETSTATIC)
                        {
                            /* static accessor takes null argument for owning object */
                            pushNull();
                        }
                        else if (opcode == PUTSTATIC)
                        {
                            methodPrefix = "set";
                            if (isCategory2)
                            {
                                /* swap cat 1 and cat 2 */
                                dupX2();
                                pop();
                                /* static accessor takes null argument for owning object */
                                pushNull();
                                /* swap cat 1 and cat 2 */
                                dupX2();
                                pop();
                            }
                            else
                            {
                                swap();
                                /* static accessor takes null argument for owning object */
                                pushNull();
                                swap();
                            }
                        }

                        /* construct accessor method type information */
                        String methodSuffix = null;
                        String fieldDesc = Type.getDescriptor(field.getType());
                        if (field.getType() == int.class)
                            methodSuffix = "Int";
                        else if (field.getType() == short.class)
                            methodSuffix = "Short";
                        else if (field.getType() == byte.class)
                            methodSuffix = "Byte";
                        else if (field.getType() == boolean.class)
                            methodSuffix = "Boolean";
                        else if (field.getType() == char.class)
                            methodSuffix = "Char";
                        else if (field.getType() == float.class)
                            methodSuffix = "Float";
                        else if (field.getType() == long.class)
                            methodSuffix = "Long";
                        else if (field.getType() == double.class)
                            methodSuffix = "Double";
                        else
                        {
                            fieldDesc = "Ljava/lang/Object;";
                            methodSuffix = "";
                        }

                        String methodSig = "(Ljava/lang/Object;";
                        if (methodPrefix.equals("get"))
                            methodSig += ")" + fieldDesc;
                        else
                            methodSig += fieldDesc + ")V";
                        /* invoke field accessor */
                        invokeVirtual("java/lang/reflect/Field", methodPrefix + methodSuffix, methodSig);

                        if (methodPrefix.equals("get") && methodSuffix.equals("") && field.getType() != Object.class)
                        {
                            /* cast retrieved value to field type */
                            checkCast(Type.getType(field.getType()));
                        }

                        return;
                    }
                }
                catch (Exception ex)
                {
                    throw undeclared(ex);
                }
            }

            mv.visitFieldInsn(opcode, owner, name, desc);
        }
    }
}
